import pickle
import numpy as np
import matplotlib.pyplot as plt


# 以开盘点作为y=kx+b的x拟合出参数k和b预测收盘点
# 线性回归模型

class LinerRegressionModel(object):

    # 定义数据初始化

    def __init__(self, data):
        self.data = data
        self.x = data[:, 0]
        self.y = data[:, 1]

    # 定义输出日志，打印开盘点和收盘点的线性关系

    def log(self, a, b):
        print("计算出的线性回归函数为:\ny = {:.5f}x + {:.5f}".format(a, b))

    # 绘图，一条拟合出的y=wx+b的线和数据点集

    def plt(self, a, b):
        plt.plot(self.x, self.y, 'o', label='data', markersize=10)
        plt.plot(self.x, a * self.x + b, 'r', label='line')
        plt.legend()
        plt.show()

    # 利用最小二乘法拟合数据点集

    def least_square_method(self):
        def calc_ab(x, y):
            sum_x, sum_y, sum_xy, sum_xx = 0, 0, 0, 0
            n = len(x)
            for i in range(0, n):
                sum_x += x[i]
                sum_y += y[i]
                sum_xy += x[i] * y[i]
                sum_xx += x[i] ** 2
            a = (sum_xy - (1 / n) * (sum_x * sum_y)) / (sum_xx - (1 / n) * sum_x ** 2)
            b = sum_y / n - a * sum_x / n
            return a, b

        a, b = calc_ab(self.x, self.y)
        # 执行日志函数
        self.log(a, b)
        # 绘图
        self.plt(a, b)
        return a


# 主函数

if __name__ == '__main__':
    # 所有的股票
    stock_choice = ['600150.XSHG', '600900.XSHG', '600048.XSHG', '600340.XSHG', '600569.XSHG',
                    '600019.XSHG', '600115.XSHG', '600118.XSHG', '600151.XSHG', '000001.XSHE',
                    '600030.XSHG', '601001.XSHG', '601857.XSHG', '600028.XSHG', '600050.XSHG',
                    '000063.XSHE', '000651.XSHE', '600111.XSHG', '600518.XSHG', '600056.XSHG',
                    '600519.XSHG']

    result = {}
    # 读取数据
    for stock in stock_choice:
        with open('data.pkl', 'rb') as f:
            data = pickle.load(f)
        stocks_close = []
        stocks_open = []
        for day in data:
            # 获取每一天的数据
            df = data[day]
            stocks_close.append(df["close"][stock])
            stocks_open.append(df["open"][stock])
        # 开盘
        a = np.array(stocks_open)
        # 收盘
        b = np.array(stocks_close)
        # 合并矩阵
        a = np.concatenate((a, b), axis=0).reshape(2431, 2)
        data = a
        # 线性回归模型
        model = LinerRegressionModel(data)
        k = model.least_square_method()
        result[stock] = k
    # 按字典的值进行排序，倒序
    sorted_key_list = sorted(result.items(), key=lambda x: x[1], reverse=True)
    # 把排序后key值提取出来
    ans = list(result.keys())
    print(ans[:5])
